#include "slepbase.h"
#include "epph.h"

void mcLogisticC(slep_matrix* A, slep_matrix* Y, double z, OPTS* opts,
			  slep_matrix* X, slep_vector *c_v, slep_vector* funVal, slep_vector* valueL){
	int k, i, j;
	slep_matrix *P_FLAG, *B;
	slep_vector *m1, *m2;
	slep_matrix *ATB, *AX;
	double lambda, q;

	k = Y->n;
	assert( Y->m == A->m);
	assert( z>0 );

	sll_opts(opts);
	initNormalization(A,opts);

	if( hasGot(opts,SlepOPT_q))
		q = opts->q = 2.0;
	else{
		q = opts->q;
		assert( SLEPDoubleApproxEquals(q,2.0));
	}

	P_FLAG = slepAllocDenseMatrix(Y->m, Y->n);
	m1 = slepAllocDenseVector(k);
	m2 = slepAllocDenseVector(k);
	B = slepAllocDenseMatrix(Y->m, Y->n);
	ATB = slepAllocDenseMatrix( X->m, X->n);
	AX = slepAllocDenseMatrix( Y->m, Y->n );
	for(j=0; j< Y->n; ++j){
		m1->data.vctr[j] = 0.0;
		for(i=0; i < Y->m; ++i){
			P_FLAG->data.mtrx[j*P_FLAG->ld+i] = (SLEPDoubleApproxEquals(Y->data.mtrx[j*Y->ld+i],1.0)?1.0:0.0);
			m1->data.vctr[j] += P_FLAG->data.mtrx[j*P_FLAG->ld+i];
		}
		m2->data.vctr[j] = Y->m - m1->data.vctr[j];
	}

	if( opts->init == 2){
		slep_zerom(X);
		slep_zerov(c_v);
	}else{
		if(hasGot(opts,SlepOPT_X0)){
			assert( opts->X0->m == X->m && opts->X0->n == X->n);
			slep_dmcopy(opts->X0, X);
		}else
			slep_zerom(X);
		if(hasGot(opts,SlepOPT_c0)){
			assert(opts->c0->dim == c_v->dim);
			slep_dvcopy(opts->c0, c_v);
		}else
			for(i=0; i< c_v->dim; ++i)
				c_v->data.vctr[i*c_v->inc] = log( m1->data.vctr[i] / m2->data.vctr[i]);
	}

	// TODO Changed it to NORMALIZEDMM
	slep_dgemm(SlepNoTranspose,SlepNoTranspose,1.0,A,X,0.0,AX);

	/* Armijo Goldstein line search scheme */
	if( opts->lFlag==0){
		int bFlag = 0, iterStep;
		double L = 1.0 / ( Y->m * Y->n );
		double alphap , alpha, beta;
		slep_matrix *Xp, *AXp, *XXp;
		slep_vector *cp, *ccp, *sc, *gc;
		slep_matrix *S, *AS;
		slep_matrix *aa, *bb, *prob;
		slep_matrix *G, *V;
		double fun_s, fun_x;
		double r_sum, l_sum;
		double lambda0;
		int steps;

		alphap = 0.0;
		alpha = 1.0;

		lambda0 = 0.0;

		XXp = slepAllocDenseMatrix(X->m, X->n);
		cp = slepAllocDenseVector(c_v->dim);
		ccp = slepAllocDenseVector(c_v->dim);
		sc = slepAllocDenseVector(c_v->dim);
		gc = slepAllocDenseVector(c_v->dim);
		S = slepAllocDenseMatrix(X->m, X->n);
		AS = slepAllocDenseMatrix(AX->m, AX->n);
		aa = slepAllocDenseMatrix(Y->m, Y->n);
		bb = slepAllocDenseMatrix(Y->m, Y->n);
		prob = slepAllocDenseMatrix(Y->m, Y->n);
		G = slepAllocDenseMatrix(X->m, X->n);
		V = slepAllocDenseMatrix(X->m, X->n);

		Xp = slepCopyToNewMatrix(X);
		AXp = slepCopyToNewMatrix(AX);
		slep_zerom(XXp);
		slep_dvcopy(c_v, cp);
		slep_zerov(ccp);

		for( iterStep = 0; iterStep < opts->maxIter; ++ iterStep ){
			beta = (alphap -1.0)/ alpha;

			slep_dmcopy(X,S);
			slep_daXpbY(SlepNoTranspose, beta, XXp, 1.0, S);

			slep_dvcopy(c_v,sc);
			slep_daxpby(beta, ccp, 1.0, sc);

			slep_dmcopy(AX, AS);
			slep_daXpbY(SlepNoTranspose, -beta, AXp, 1.0+beta, AS);

			fun_s = 0.0;
			for(j=0; j< Y->n; ++j){
				gc->data.vctr[j] = 0.0;
				for(i=0; i< Y->m; ++i){
					aa->data.mtrx[j*aa->ld+i] = -Y->data.mtrx[j*Y->ld+i] *( AS->data.mtrx[j*AS->ld+i] + sc->data.vctr[j] );
					bb->data.mtrx[j*bb->ld+i] = SLEPMAX(aa->data.mtrx[j*aa->ld+i], 0.0);
					fun_s += log( exp(- bb->data.mtrx[j*bb->ld+i]) + exp(aa->data.mtrx[j*aa->ld+i] - bb->data.mtrx[j*bb->ld+i])) + bb->data.mtrx[j*bb->ld+i] ;
					prob->data.mtrx[j*prob->ld+i] = 1.0/(1.0 + exp(aa->data.mtrx[j*aa->ld+i]));
					B->data.mtrx[j*B->ld+i] = -Y->data.mtrx[j*Y->ld+i] *(1.0 - prob->data.mtrx[j*prob->ld+i]) / (prob->m * prob->n);
					gc->data.vctr[j] += B->data.mtrx[j*B->ld+i];
				}
			}
			fun_s /= (Y->m * Y->n);

			// G = ATB
			slep_dgemm(SlepTranspose,SlepNoTranspose,1.0,A,B,0.0,G);
			
			slep_dmcopy(X,Xp);
			slep_dmcopy(AX,AXp);
			slep_dvcopy(c_v, cp);
			for(;;){
				slep_dmcopy(S,V);
				slep_daXpbY(SlepNoTranspose, -1.0/L, G, 1.0, V);
				slep_dvcopy(sc,c_v);
				slep_daxpby(-1.0/L,gc,1.0,c_v);

				ep21d(X->data.mtrx,&lambda,&steps,V->data.mtrx,X->m, X->n, z, lambda0);
				lambda0 = lambda;

				slep_dmcopy(X, V);
				slep_daXpbY(SlepNoTranspose, -1.0, S, 1.0, V);

				slep_dgemm(SlepNoTranspose, SlepNoTranspose, 1.0, A, X, 0.0 , AX);
				fun_x = 0.0;
				for(j=0; j< Y->n; ++j)
					for(i=0; i< Y->m; ++i){
						aa->data.mtrx[j*aa->ld+i] = -Y->data.mtrx[j*Y->ld+i] *( AX->data.mtrx[j*AX->ld+i] + c_v->data.vctr[j] );
						bb->data.mtrx[j*bb->ld+i] = SLEPMAX(aa->data.mtrx[j*aa->ld+i], 0.0);
						fun_x += log( exp(- bb->data.mtrx[j*bb->ld+i]) + exp(aa->data.mtrx[j*aa->ld+i] - bb->data.mtrx[j*bb->ld+i])) + bb->data.mtrx[j*bb->ld+i] ;
					}
				fun_x /= (Y->m * Y->n);
				r_sum = slep_dfrbnrm(V);
				r_sum = r_sum * r_sum;
				r_sum += slep_ddot(c_v,c_v) + slep_ddot(sc,sc) - 2.0 * slep_ddot(sc,c_v);
				r_sum /= 2.0;
				l_sum = fun_x - fun_s - slep_dtraceAB(SlepNoTranspose,V, SlepTranspose,G) - slep_ddot(c_v,gc) + slep_ddot(sc,gc);
				if( r_sum <= 1e-20 ){
					bFlag = 1;
					break;
				}
				if( l_sum <= r_sum * L)
					break;
				else
					L = SLEPMAX(2.0*L, l_sum/r_sum);
			}
			alphap = alpha;
			alpha = (1.0 + sqrt(4.0* alpha *alpha + 1.0))/2.0;

			valueL->data.vctr[iterStep] = L ;
			slep_dmcopy(X,XXp);
			slep_daXpbY(SlepNoTranspose, -1.0, Xp, 1.0, XXp);
			slep_dvcopy(c_v, ccp);
			slep_daxpby(-1.0,cp,1.0,ccp);
			funVal->data.vctr[iterStep] = fun_x;

			if(bFlag)
				break;
			if( opts->tFlag == 0){
				if( iterStep >= 2){
					if( fabs(funVal->data.vctr[iterStep] - funVal->data.vctr[iterStep-1]) < opts->tol)
						break;
				}
			}else if( opts->tFlag == 1){
				if(iterStep>=2){
					if( fabs(funVal->data.vctr[iterStep] - funVal->data.vctr[iterStep-1]) < (opts->tol)*(funVal->data.vctr[iterStep-1]) )
						break;
				}
			}else if(opts->tFlag == 2){
				if( funVal->data.vctr[iterStep] <= opts->tol)
					break;
			}else if(opts->tFlag == 3){
				if(slep_dfrbnrm(XXp) <= opts->tol)
					break;
			}else if(opts->tFlag == 4){
				if( slep_dfrbnrm(XXp) <= (opts->tol) * max(slep_dfrbnrm(Xp),1))
					break;
			}else if(opts->tFlag == 5){
				if( iterStep >= opts->maxIter)
					break;
			}
		}
		slepFreeMatrix(Xp);
		slepFreeMatrix(AXp);
		slepFreeMatrix(XXp);

		slepFreeVector(cp);
		slepFreeVector(ccp);
		slepFreeVector(sc);
		slepFreeVector(gc);

		slepFreeMatrix(S);
		slepFreeMatrix(AS);
		slepFreeMatrix(aa);
		slepFreeMatrix(bb);
		slepFreeMatrix(prob);
		slepFreeMatrix(G);
		slepFreeMatrix(V);
		funVal->dim = iterStep;
		valueL->dim = iterStep;

	}

	/* Adaptive Line search*/
	if( opts->lFlag==1){
		int bFlag = 0, iterStep;
		double L = 1.0 / ( Y->m * Y->n );
		double alphap, alpha, beta;
		slep_matrix *Xp, *AXp, *XXp;
		slep_vector *cp, *ccp, *sc, *gc;
		slep_matrix *S, *AS;
		slep_matrix *aa, *bb, *prob;
		slep_matrix *G, *V;
		double fun_s, fun_x;
		double r_sum, l_sum;
		double gamma, tao;
		double lambda0 ;
		int steps;

		slep_vector *cnew;
		slep_matrix *Xnew, *AXnew;


		gamma = 1.0;
		lambda0 = 0.0;

		XXp = slepAllocDenseMatrix(X->m, X->n);
		cp = slepAllocDenseVector(c_v->dim);
		ccp = slepAllocDenseVector(c_v->dim);
		sc = slepAllocDenseVector(c_v->dim);
		gc = slepAllocDenseVector(c_v->dim);
		S = slepAllocDenseMatrix(X->m, X->n);
		AS = slepAllocDenseMatrix(AX->m, AX->n);
		aa = slepAllocDenseMatrix(Y->m, Y->n);
		bb = slepAllocDenseMatrix(Y->m, Y->n);
		prob = slepAllocDenseMatrix(Y->m, Y->n);
		G = slepAllocDenseMatrix(X->m, X->n);
		V = slepAllocDenseMatrix(X->m, X->n);

		Xnew = slepAllocDenseMatrix(X->m, X->n);
		AXnew = slepAllocDenseMatrix(AX->m, AX->n);

		cnew = slepAllocDenseVector(c_v->dim);

		Xp = slepCopyToNewMatrix(X);
		AXp = slepCopyToNewMatrix(AX);
		slep_zerom(XXp);
		slep_dvcopy(c_v, cp);
		slep_zerov(ccp);

		for( iterStep = 0; iterStep < opts->maxIter; ++ iterStep ){
			for(;;){
				if(iterStep != 0){
					alpha = (-gamma + sqrt(gamma *gamma + 4.0 * L * gamma ))/(2.0*L);
					beta = (gamma - gamma * alphap) /(alphap *gamma + alphap * L * alpha);
					slep_dmcopy(X,S);
					slep_daXpbY(SlepNoTranspose, beta, XXp, 1.0, S);

					slep_dvcopy(c_v,sc);
					slep_daxpby(beta, ccp, 1.0, sc);

					slep_dmcopy(AX, AS);
					slep_daXpbY(SlepNoTranspose, -beta, AXp, 1.0+beta, AS);
				}else{
					alpha = (-1.0 + sqrt(5.0)) / 2.0;
					beta = 0.0;
					slep_dmcopy(X,S);
					slep_dvcopy(c_v,sc);
					slep_dmcopy(AX,AS);
				}


				fun_s = 0.0;
				for(j=0; j< Y->n; ++j){
					gc->data.vctr[j] = 0.0;
					for(i=0; i< Y->m; ++i){
						aa->data.mtrx[j*aa->ld+i] = -Y->data.mtrx[j*Y->ld+i] *( AS->data.mtrx[j*AS->ld+i] + sc->data.vctr[j] );
						bb->data.mtrx[j*bb->ld+i] = SLEPMAX(aa->data.mtrx[j*aa->ld+i], 0.0);
						fun_s += log( exp(- bb->data.mtrx[j*bb->ld+i]) + exp(aa->data.mtrx[j*aa->ld+i] - bb->data.mtrx[j*bb->ld+i])) + bb->data.mtrx[j*bb->ld+i] ;
						prob->data.mtrx[j*prob->ld+i] = 1.0/(1.0 + exp(aa->data.mtrx[j*aa->ld+i]));
						B->data.mtrx[j*B->ld+i] = -Y->data.mtrx[j*Y->ld+i] *(1.0 - prob->data.mtrx[j*prob->ld+i]) / (prob->m * prob->n);
						gc->data.vctr[j] += B->data.mtrx[j*B->ld+i];
					}
				}
				fun_s /= (Y->m * Y->n);

				// G = ATB
				slep_dgemm(SlepTranspose,SlepNoTranspose,1.0,A,B,0.0,G);
			


				slep_dmcopy(S,V);
				slep_daXpbY(SlepNoTranspose, -1.0/L, G, 1.0, V);
				slep_dvcopy(sc,cnew);
				slep_daxpby(-1.0/L,gc,1.0,cnew);

				ep21d(Xnew->data.mtrx,&lambda,&steps,V->data.mtrx,X->m,X->n,z,lambda0);
				lambda0 = lambda;

				slep_dmcopy(Xnew, V);
				slep_daXpbY(SlepNoTranspose, -1.0, S, 1.0, V);

				slep_dgemm(SlepNoTranspose, SlepNoTranspose, 1.0, A, Xnew, 0.0 , AXnew);
				fun_x = 0.0;
				for(j=0; j< Y->n; ++j)
					for(i=0; i< Y->m; ++i){
						aa->data.mtrx[j*aa->ld+i] = -Y->data.mtrx[j*Y->ld+i] *( AXnew->data.mtrx[j*AXnew->ld+i] + cnew->data.vctr[j] );
						bb->data.mtrx[j*bb->ld+i] = SLEPMAX(aa->data.mtrx[j*aa->ld+i], 0.0);
						fun_x += log( exp(- bb->data.mtrx[j*bb->ld+i]) + exp(aa->data.mtrx[j*aa->ld+i] - bb->data.mtrx[j*bb->ld+i])) + bb->data.mtrx[j*bb->ld+i] ;
					}
				fun_x /= (Y->m * Y->n);
				r_sum = slep_dfrbnrm(V);
				r_sum = r_sum * r_sum;
				r_sum += slep_ddot(cnew,cnew) + slep_ddot(sc,sc) - 2.0 * slep_ddot(sc,cnew);
				r_sum /= 2.0;
				l_sum = fun_x - fun_s - slep_dtraceAB(SlepNoTranspose,V, SlepTranspose,G) - slep_ddot(cnew,gc) + slep_ddot(sc,gc);
				if( r_sum <= 1e-20 ){
					bFlag = 1;
					break;
				}
				if( l_sum <= r_sum * L)
					break;
				else
					L = SLEPMAX(2.0*L, l_sum/r_sum);
			}
			gamma = L * alpha * alpha;
			alphap = alpha;
			valueL->data.vctr[iterStep] = L ;

			tao = L * r_sum/ l_sum;
			if(tao >= 5.0)
				L = L *0.8;
			slep_dmcopy(X,Xp);
			slep_dmcopy(Xnew,X);
			slep_dmcopy(AX,AXp);
			slep_dmcopy(AXnew,AX);
			slep_dvcopy(c_v, cp);
			slep_dvcopy(cnew, c_v);

			slep_dmcopy(X,XXp);
			slep_daXpbY(SlepNoTranspose, -1.0, Xp, 1.0, XXp);
			slep_dvcopy(c_v, ccp);
			slep_daxpby(-1.0,cp,1.0,ccp);
			funVal->data.vctr[iterStep] = fun_x ;
			if(bFlag)
				break;
			if( opts->tFlag == 0){
				if( iterStep >= 2){
					if( fabs(funVal->data.vctr[iterStep] - funVal->data.vctr[iterStep-1]) < opts->tol)
						break;
				}
			}else if( opts->tFlag == 1){
				if(iterStep>=2){
					if( fabs(funVal->data.vctr[iterStep] - funVal->data.vctr[iterStep-1]) < (opts->tol)*(funVal->data.vctr[iterStep-1]) )
						break;
				}
			}else if(opts->tFlag == 2){
				if( funVal->data.vctr[iterStep] <= opts->tol)
					break;
			}else if(opts->tFlag == 3){
				if(slep_dfrbnrm(XXp) <= opts->tol)
					break;
			}else if(opts->tFlag == 4){
				if( slep_dfrbnrm(XXp) <= (opts->tol) * max(slep_dfrbnrm(Xp),1))
					break;
			}else if(opts->tFlag == 5){
				if( iterStep >= opts->maxIter)
					break;
			}
		}
		slepFreeMatrix(Xp);
		slepFreeMatrix(AXp);
		slepFreeMatrix(XXp);

		slepFreeVector(cp);
		slepFreeVector(ccp);
		slepFreeVector(sc);
		slepFreeVector(gc);

		slepFreeMatrix(S);
		slepFreeMatrix(AS);
		slepFreeMatrix(aa);
		slepFreeMatrix(bb);
		slepFreeMatrix(prob);
		slepFreeMatrix(G);
		slepFreeMatrix(V);

		slepFreeMatrix(Xnew);
		slepFreeMatrix(AXnew);
		slepFreeVector(cnew);
		funVal->dim = iterStep;
		valueL->dim = iterStep;

	}

	slepFreeMatrix(P_FLAG);
	slepFreeMatrix(B);
	slepFreeVector(m1);
	slepFreeVector(m2);
	slepFreeMatrix(ATB);
	slepFreeMatrix(AX);
}